package com.xdh.springcloud.security.jwt.auth.ajax;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.xdh.springcloud.common.message.Message;
import com.xdh.springcloud.common.message.MessageCode;
import com.xdh.springcloud.security.jwt.common.factory.JwtTokenFactory;
import com.xdh.springcloud.security.jwt.common.model.JwtToken;
import com.xdh.springcloud.security.jwt.common.model.UserModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.security.core.Authentication;
import org.springframework.security.web.WebAttributes;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;

@Component
public class AjaxAwareAuthenticationSuccessHandler implements AuthenticationSuccessHandler {

    @Autowired
    private ObjectMapper objectMapper;
    @Autowired
    private JwtTokenFactory tokenFactory;

    @Override
    public void onAuthenticationSuccess(HttpServletRequest request, HttpServletResponse response, Authentication authentication) throws IOException {
        UserModel userModel = (UserModel) authentication.getPrincipal();
        JwtToken jwtToken = tokenFactory.createAccessJwtToken(userModel);
        Message message = tokenFactory.createMessage(jwtToken);
        message.setCode(MessageCode.SUCCESS.getCode());
        response.setStatus(HttpStatus.OK.value());
        response.setContentType(MediaType.APPLICATION_JSON_VALUE);
        objectMapper.writeValue(response.getWriter(), message);
        clearAuthenticationAttributes(request);
    }

    /**
     * Removes temporary authentication-related data which may have been stored
     * in the session during the authentication process..
     */
    protected final void clearAuthenticationAttributes(HttpServletRequest request) {
        HttpSession session = request.getSession(false);
        if (session == null) {
            return;
        }
        session.removeAttribute(WebAttributes.AUTHENTICATION_EXCEPTION);
    }
}
